import os
import json
import ast
import hydra
from hydra import utils
import logging
from easyinstruct.prompts import IEPrompt
from preprocess import prepare_examples
import csv
from chatgpt import ask_anything

logger = logging.getLogger(__name__)

RTE_CH_DOMAIN = "您是一个高度智能和精确的{0}域资源描述框架（RDF）数据模型。您将文本作为输入，并将其转换为{0}域RDF三元组。三元组是由三个实体组成的集合，以主语-谓语-宾语表达式的形式对语义数据进行编码。\n"
# RTE_CH_REL = ("您是一个拥有丰富医学领域知识的专家，完成以下任务。"
#           "将输入文本转换为三元组格式的数据，三元组是由三个词语组成的集合，以主语-谓语-宾语表达式的形式对语义数据进行编码。要求写出尽可能多的的三元组，主谓宾都是简短的词语，不要重复。您输出的格式需要为"
#           "[[主语类型, 主语, 宾语类型, 宾语, 关系],...]"
#           "其中，主语类型只有检查项目、科室、疾病、药品、食物、症状这6种，宾语类型只有检查项目、科室、疾病、药品、食物、症状这6种，关系只有‘疾病属于科室’、‘疾病常用药品’、‘疾病推荐药品’、‘疾病宜吃食物’、‘疾病推荐食谱’、‘疾病忌吃食物’、‘疾病所需检查’、‘疾病症状’、‘疾病并发症’这9种。"
#           "第一排是标签，后面是数据，要换行，字符串请用双引号括起来，确保主语和宾语都不是动作，除了以上格式外不要输出任何其他的字符。\n\n")
RTE_CH_ATT = ("您是一个拥有丰富医学领域知识的专家，完成以下任务。"
          "将输入文本转换为三元组格式的数据，三元组是由三个部分组成的集合，以主语-属性类型-属性值的形式对语义数据进行编码。要求写出尽可能多的的三元组，主语和谓语都是简短的词语，宾语可以是词语或者一段话。同一个主语和谓语只保留一个三元组。您输出的格式需要为"
          "[[主语, 属性类型, 属性值],...]"
          "其中，主语必须是某种疾病，属性类型只有疾病名称、疾病简介、疾病病因、预防措施、治疗周期、治疗方式、治愈概率、疾病易感人群这8种。"
          "数据要换行，字符串请用双引号括起来，确保主语和宾语都不是动作，除了以上格式外不要输出任何其他的字符。\n\n")
RTE_CH_OUTPUT = "您输出的格式需要为[[(主语)-[:谓语]->(宾语)], ...]，没有其他格式要求。\n\n"

"""
三元组的构建
"""

@hydra.main(version_base=None, config_path="", config_name="run")
def main(cfg):
    print("Hello!")
    cfg.cwd = utils.get_original_cwd()

    text = cfg.text_input

    if not cfg.api_key:
        raise ValueError("Need an API Key.")
    if cfg.engine not in ["text-davinci-003", "text-davinci-002", "text-curie-001", "text-babbage-001", "text-ada-001", "gpt-3.5-turbo", "gpt-4"]:
        raise ValueError("The OpenAI model is not supported now.")

    os.environ['OPENAI_API_KEY'] = cfg.api_key
    os.environ["OPENAI_BASE_URL"] = cfg.base_url

    # # 指定文件路径
    # input_file_path = 'example.txt'
    #
    # # 打开文件
    # with open(input_file_path, 'r', encoding='utf-8') as file:
    #     # 读取文件全部内容到字符串
    #     content = file.read()
    #
    # # 需要读取输入的文件来完成三元组抽取任务。
    #
    #
    # # 打印结果字符串
    # print(content)
    # region get_relation
    # cfg.instruction = RTE_CH_REL
    cfg.instruction = RTE_CH_ATT
    ie_prompter = IEPrompt(cfg.task)

    examples = None
    if not cfg.zero_shot:
        examples = prepare_examples(cfg.data_path, cfg.task, cfg.language)

    if cfg.task == 're':
        ie_prompter.build_prompt(
            prompt=text,
            head_entity=cfg.head_entity,
            head_type=cfg.head_type,
            tail_entity=cfg.tail_entity,
            tail_type=cfg.tail_type,
            language=cfg.language,
            instruction=cfg.instruction,
            in_context=not cfg.zero_shot,
            domain=cfg.domain,
            labels=cfg.labels,
            examples=examples
        )
    else:
        ie_prompter.build_prompt(
            prompt=text,
            language=cfg.language,
            instruction=cfg.instruction,
            in_context=not cfg.zero_shot,
            domain=cfg.domain,
            labels=cfg.labels,
            examples=examples
        )

    # result = ie_prompter.get_openai_result()
    result = ask_anything(RTE_CH_ATT + text)
    logger.info(result)

    # 将字符串转换为Python列表
    # data_list = eval(res)
    data_list = ast.literal_eval(result)
    print("data list:",data_list)

    # 打开文件用于写入
    with open('data_att.csv', 'w', newline='', encoding='utf-8') as csvfile:
        writer = csv.writer(csvfile)

        # 写入数据
        for row in data_list:
            writer.writerow(row)

    print("实体-属性数据已写入CSV文件。")
    # endregion

    # cfg.instruction = RTE_CH_ATT
    # ie_prompter = IEPrompt(cfg.task)
    #
    # examples = None
    # if not cfg.zero_shot:
    #     examples = prepare_examples(cfg.data_path, cfg.task, cfg.language)
    #
    # if cfg.task == 're':
    #     ie_prompter.build_prompt(
    #         prompt=text,
    #         head_entity=cfg.head_entity,
    #         head_type=cfg.head_type,
    #         tail_entity=cfg.tail_entity,
    #         tail_type=cfg.tail_type,
    #         language=cfg.language,
    #         instruction=cfg.instruction,
    #         in_context=not cfg.zero_shot,
    #         domain=cfg.domain,
    #         labels=cfg.labels,
    #         examples=examples
    #     )
    # else:
    #     ie_prompter.build_prompt(
    #         prompt=text,
    #         language=cfg.language,
    #         instruction=cfg.instruction,
    #         in_context=not cfg.zero_shot,
    #         domain=cfg.domain,
    #         labels=cfg.labels,
    #         examples=examples
    #     )
    #
    # result = ie_prompter.get_openai_result()
    # res = "[[head_type, head_entity, tail_type, tail_entity, relation],[Event, 2019年美国网球公开赛男单冠军, Champion, 亚瑟·阿什球场, 获得],[Person, 纳达尔, Event, 2019年美国网球公开赛男单冠军, 参加],[Person, 纳达尔, Opponent, 梅德韦杰夫, 击败],[Opponent, 梅德韦杰夫, Seed, 5, 是],[Time, 4小时51分钟, Event, 2019年美国网球公开赛男单冠军, 历时],[Time, 北京时间9月9日上午, Event, 2019年美国网球公开赛男单冠军, 发生]]"
    # logger.info(result)
    #
    # # 将字符串转换为Python列表
    # # data_list = eval(res)
    # data_list = ast.literal_eval(result)
    # print("data list:", data_list)
    #
    # # 打开文件用于写入
    # with open('data.csv', 'a', newline='', encoding='utf-8') as csvfile:
    #     writer = csv.writer(csvfile)
    #
    #     # 写入数据
    #     for row in data_list:
    #         writer.writerow(row)
    #
    # print("实体-属性数据已写入CSV文件。")
    # endregion

if __name__ == '__main__':
    main()